import argparse
import logging
import os
import platform
import pdb
import math
import sys
import time

import numpy as np
from matplotlib import animation
import matplotlib.pyplot as plt
import matplotlib.patches
from sklearn.preprocessing import normalize

# try:
#     
# except IndexError:
#     print("Cannot add the common path {}".format(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# print(sys.path)
from src.motionplanner.Constrained_ILQR.scripts.arguments import add_arguments
from PolyRect import PolyRect
from src.motionplanner.Constrained_ILQR.scripts.ilqr.iLQR import iLQR

# Add lambda functions
cos = lambda a : np.cos(a)
sin = lambda a : np.sin(a)
tan = lambda a : np.tan(a)

PI = math.pi
colors = ['r', 'g', 'b', 'k']

class SimParams:
    dt = 0.1               # 仿真时间步长
    sim_time = 1         # 总仿真步数
    map_lengthx = 100       # 地图在 X 方向的长度（用于画图）
    map_lengthy = 6.8        # 地图在 Y 方向的半高度（用于画图）
    lane1 = 3.4
    lane2 = 0
    lane3 = -3.4
    num_vehicles = 3       # 车辆总数

    # 车辆相关参数
    car_dims = np.array([1.4, 0.94])           # 车辆长和宽
    start_state = np.array([46, - 3.4 / 2, 0, 0])  # Ego 车的初始状态 [x, y, 速度v, 航向角theta]
    max_speed = 180/3.6                  # 最大速度 (180km/h -> m/s)
    wheelbase = 1.18                     # 轴距
    steer_min = -1.0                     # 最小转向角（弧度）
    steer_max = 1.0                      # 最大转向角（弧度）
    accel_min = -5.5                     # 最小加速度
    accel_max = 5                      # 最大加速度
    desired_y = - 3.4 / 2                      # 期望行驶的车道对应的 y 坐标
    NPC_max_acc = 0.75                   # NPC车辆的最大加速度

# Lanes defined at 4, 0, -4

class PySimulator:

    def __init__(self, args, SimParams, NPC_start, NPC_control):
        """
        :param args:          命令行解析的参数
        :param SimParams:     仿真全局配置，如时间步长/车辆尺寸等
        :param NPC_start:     NPC车辆起始状态 [x, y, v, theta]
        :param NPC_control:   NPC车辆的控制输入序列 (加速度、角速度等)
        """
        self.args = args
        # num_vehicles is only 2 for now
        # 用于存储各车辆对应的 PolyRect（矩形外形）
        self.NPC_dict = {}   # 字典 key=车辆ID, value=PolyRect实例
        self.patches = []    # 用于保存matplotlib的 Patch，用来在动画中更新位置

        self.simparams = SimParams
        self.num_vehicles = self.simparams.num_vehicles
        self.current_ego_state = self.simparams.start_state
        # self.last_ego_states = [self.simparams.start_state[0:2]]
        self.last_ego_states = np.array([self.simparams.start_state[0], self.simparams.start_state[1]])

        # Plot parameters
        # 绘图相关设置
        self.fig = plt.figure(figsize=(50, 5))
        self.ax = self.fig.add_subplot(111)
        self.ax.axis('equal')
        self.ax.set_xlim(-5, self.simparams.map_lengthx+5)
        self.ax.set_ylim(-self.simparams.map_lengthy, self.simparams.map_lengthy)
        # 绘制车道：lane1 和 lane3 为实线，lane2 为虚线
        self.ax.axhline(y=self.simparams.lane1, c='k', lw='4')
        self.ax.axhline(y=self.simparams.lane2, c='k', lw='2', ls='--')
        self.ax.axhline(y=self.simparams.lane3, c='k', lw='4')

        # Plot Local Plan
        # 用于绘制局部规划和期望规划轨迹的点
        self.x_local_plan = []
        self.y_local_plan = []
        self.x_desired_plan = []
        self.y_desired_plan = []
        self.local_plan_plot, = plt.plot([], [], 'go', ms=10)
        self.desired_plan_plot, = plt.plot([], [], 'co', ms=5)
        self.last_states_plots, = plt.plot([], [], 'mo', ms=5)
        trajs = [self.local_plan_plot, self.desired_plan_plot, self.last_states_plots,] # C3
        self.line_plots = list(trajs)


        # 初始化 iLQR agent (规划器)
        self.navigation_agent = None
        self.create_ilqr_agent()# 创建 iLQR 对象
        self.NPC_states = [] # NPC车辆的状态序列
        # 并构建 NPC 的轨迹
        self.simulate_npc(NPC_start, NPC_control)
        self.count = 0# 用来索引 NPC 状态的计数器

        # Ego vehicle is first
        # -------------------------
        # 为所有车辆分配 PolyRect 并创建对应的 Patch
        #   这里默认第 0 号是 Ego 车，其它是 NPC
        # -------------------------
        for i in range(num_vehicles):
            # 初始化每台车的多边形
            self.NPC_dict[i] = PolyRect(self.simparams.car_dims)
            # 从多边形获取一个 matplotlib Patch，用于后续更新绘制
            self.patches.append(self.NPC_dict[i].getPatch(self.ax, colors[i]))
            self.ax.add_patch(self.patches[i]) # C1
        
        # May need to add loop here
        # self.ax.add_patch(self.patches[0]) C2
        # self.ax.add_patch(self.patches[1])

        # 最后把车辆和轨迹的所有 patch 合并，以便动画函数统一返回
        self.all_patches = self.patches + self.line_plots

    def simulate_npc(self, init_state, control):
        self.NPC_states.append(init_state)
        control = np.hstack((control, np.zeros((2, self.args.horizon))))
        for i in range(control.shape[1]):
            NPC_next_state = self.run_model_simulation(self.NPC_states[i], control[:, i])
            self.NPC_states.append(NPC_next_state)
        self.NPC_states = np.array(self.NPC_states).T
        print("self.NPC_states", self.NPC_states)
        print("size", len(self.NPC_states))
        print("size2", len(self.NPC_states[0]))
        

    def create_global_plan(self):
        # -------------------------
        # 创建并保存全局计划轨迹(简单地在 y=desired_y 处画一条横线)
        # -------------------------
        y = self.simparams.desired_y
        self.plan_ilqr = []
        for i in range(0, self.simparams.map_lengthx+20):
            self.plan_ilqr.append(np.array([i, y]))
        self.plan_ilqr = np.array(self.plan_ilqr)
        self.ax.axhline(y=y, c='r', lw='2')

    # def init_sim(self):
    #     return self.patches[0], self.patches[1], self.local_plan_plot, self.desired_plan_plot,
    def init_sim(self): #C4
        return self.all_patches

    def get_ego_states(self):
        ego_states = np.array([[self.current_ego_state[0], self.current_ego_state[1],                         0],
                               [self.current_ego_state[2],                         0,                         0],
                               [                        0,                         0, self.current_ego_state[3]],
                               [                        0,                         0,                         0],
                               [                        0,                         0,                         0]])
        return ego_states

    def get_npc_bounding_box(self):
        return self.simparams.car_dims

    def get_npc_states(self, i):
        return self.NPC_states[:, i:i+self.args.horizon]

    def create_ilqr_agent(self):
        # -------------------------
        # 创建 iLQR 规划器
        # -------------------------
        self.create_global_plan()
        self.navigation_agent = iLQR(self.args, self.get_npc_bounding_box())
        # self.plan_ilqr是参考轨迹
        self.navigation_agent.set_global_plan(self.plan_ilqr)

    def run_step_ilqr(self):
        assert self.navigation_agent != None, "Navigation Agent not initialized"

        start = time.process_time()
        desired_path, local_plan, control = self.navigation_agent.run_step(self.get_ego_states(), self.get_npc_states(self.count))
        print(time.process_time() - start)

        print("Controller: Acc {} Steer: {}".format(control[0, 0], control[1, 0]))

        return desired_path, local_plan, control[:, 0]
 
    def animate(self, i):
        # Get new ego patch
        desired_path, local_plan, control = self.run_step_ilqr()
        self.current_ego_state = self.run_model_simulation(self.current_ego_state, control)

        self.last_ego_states = np.vstack((self.last_ego_states, self.current_ego_state[0:2]))
        
        self.NPC_dict[0].createCuboid([self.current_ego_state[0], self.current_ego_state[1], self.current_ego_state[3]]) # Update ego vehicle patch
        self.patches[0].set_xy(self.NPC_dict[0].getCorners()) # Update ego vehicle patch

        # Get new NPC patch
        # pdb.set_trace()
        for j in range(1,self.num_vehicles):
            self.NPC_dict[j].createCuboid([self.NPC_states[0, i], self.NPC_states[1, i], self.NPC_states[3, i]])
            self.patches[j].set_xy(self.NPC_dict[j].getCorners())

        # Get local plan
        self.x_local_plan = local_plan[:, 0]
        self.y_local_plan = local_plan[:, 1]
        self.local_plan_plot.set_data(self.x_local_plan, self.y_local_plan)
        
        self.last_states_plots.set_data(self.last_ego_states[:, 0],self.last_ego_states[:, 1])

        #Get desired plan
        self.x_desired_plan = desired_path[:, 0]
        self.y_desired_plan = desired_path[:, 1]
        self.desired_plan_plot.set_data(self.x_desired_plan, self.y_desired_plan)

        # return self.patches[0], self.patches[1], self.local_plan_plot, self.desired_plan_plot,
        return self.all_patches

    def run_simulation(self):
        anim = animation.FuncAnimation(self.fig, self.animate,
                               init_func=self.init_sim,
                               frames=self.simparams.sim_time,
                               interval=1000,
                               blit=True,
                               repeat=True)
        anim.save("./results_animation.gif", writer='ffmpeg', fps=10)
        plt.show()

    def run_model_simulation(self, state, control):
        """
        Find the next state of the vehicle given the current state and control input
        """
        # Clips the controller values between min and max accel and steer values
        control[0] = np.clip(control[0], self.simparams.accel_min, self.simparams.accel_max)
        control[1] = np.clip(control[1], state[2]*tan(self.simparams.steer_min)/self.simparams.wheelbase, state[2]*tan(self.simparams.steer_max)/self.simparams.wheelbase)
        
        Ts = self.simparams.dt
        next_state = np.array([state[0] + cos(state[3])*(state[2]*Ts + (control[0]*Ts**2)/2),
                               state[1] + sin(state[3])*(state[2]*Ts + (control[0]*Ts**2)/2),
                               np.clip(state[2] + control[0]*Ts, 0.0, self.simparams.max_speed),
                              (state[3] + control[1]*Ts)%(2*np.pi)])  # wrap angles between 0 and 2*pi
        # print("Next state {}".format(next_state))

        return next_state






if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description='CARLA CILQR')
    add_arguments(argparser)
    args = argparser.parse_args()

    NPC_start = np.array([2, 2.5, 0, 0])
    print(NPC_start)
    NPC_ctrl = np.zeros((2, SimParams.sim_time))

    NPC_ctrl[0,:] = np.linspace(SimParams.NPC_max_acc, 0, SimParams.sim_time)
    num_vehicles = 3
    print("NPC_control", NPC_ctrl)
    pysim = PySimulator(args, SimParams, NPC_start, NPC_ctrl)
    pysim.run_simulation()
